import os
import pandas as pd
import numpy as np
from pysr import PySRRegressor
import time
from bishe_situations.utils import BINARY_OPS_PYSR, UNARY_OPS_PYSR
import json

def run_pysr_experiment():
    # 创建结果保存目录
    results_dir = "pysr_results"
    os.makedirs(results_dir, exist_ok=True)
    
    # 加载参数变体
    with open('param_variants/pysr_variants.json', 'r') as f:
        param_data = json.load(f)
        param_variants = param_data['variants']
    
    # 获取所有数据文件
    data_dir = "nguyen_data"
    data_files = [f for f in os.listdir(data_dir) if f.endswith('_data.csv')]
    
    # 存储所有结果
    all_results = []
    
    for data_file in data_files:
        print(f"\n处理文件: {data_file}")
        equation_name = data_file.replace('_data.csv', '')
        
        # 读取数据
        df = pd.read_csv(os.path.join(data_dir, data_file))
        X = df['X'].values.reshape(-1, 1)
        y = df['y'].values
        
        # 为每个参数变体运行实验
        for i, variant_params in enumerate(param_variants):
            print(f"运行参数变体 {i+1}/{len(param_variants)}")
            
            # 基础参数
            params = {}
            # 更新变体参数
            params.update(variant_params)
            
            # 记录开始时间
            start_time = time.time()
            
            # 训练模型
            model = PySRRegressor(                
                binary_operators=BINARY_OPS_PYSR,
                unary_operators=UNARY_OPS_PYSR,
                niterations=20,
                population_size=20,
                maxsize=20,
                parsimony=0.1,
                progress=False,
                verbosity=1,
                constraints={'^': (-5, 5)},
                **params
            )
            model.fit(X, y)
            
            # 计算训练时间
            train_time = time.time() - start_time
            
            # 获取最佳方程和MSE
            best_equation = model.sympy()
            best_mse = model.score(X, y)
            
            # 保存结果
            result = {
                'equation_name': equation_name,
                'variant_index': i,
                'best_equation': str(best_equation),
                'mse': float(best_mse),
                'train_time': train_time,
                'params': params
            }
            all_results.append(result)
            
            # 打印结果
            print(f"方程: {equation_name}")
            print(f"变体: {i+1}")
            print(f"最佳方程: {best_equation}")
            print(f"MSE: {best_mse:.6f}")
            print(f"训练时间: {train_time:.2f}秒")
            
            # 每100个变体保存一次结果
            if (i + 1) % 100 == 0:
                # 保存当前结果
                current_results_file = os.path.join(results_dir, f"{equation_name}_results_{i+1}.json")
                with open(current_results_file, 'w') as f:
                    json.dump(all_results, f, indent=4)
                
                # 创建当前汇总表格
                summary = pd.DataFrame(all_results)
                summary_file = os.path.join(results_dir, f"{equation_name}_summary_{i+1}.csv")
                summary.to_csv(summary_file, index=False)
    
    # 保存所有结果
    all_results_file = os.path.join(results_dir, "all_results.json")
    with open(all_results_file, 'w') as f:
        json.dump(all_results, f, indent=4)
    
    # 创建最终汇总表格
    summary = pd.DataFrame(all_results)
    summary_file = os.path.join(results_dir, "summary.csv")
    summary.to_csv(summary_file, index=False)
    
    print("\n所有结果已保存到:", results_dir)

if __name__ == "__main__":
    run_pysr_experiment() 